import decimal
import math
import queue
import sys
from datetime import datetime

import pandas as pd
import psutil


def is_port_listening(port):
    for conn in psutil.net_connections(kind='inet'):
        if conn.laddr.port == port and conn.status == 'LISTEN':
            return True
    return False


def format_timestamp(param):
    param = int(param)
    dt_object = datetime.fromtimestamp(param / 1000)
    return dt_object.strftime("%Y-%m-%d %H:%M:%S.") + str(param % 1000).zfill(3)


def split_data(origin_list, is_df=False, use_protocol=True):
    if not is_df:
        origin_list = pd.DataFrame(origin_list)

    # normal和webshell都只要包含http就可以预测
    normal_df = origin_list[
        # origin_list['protocol'].str.contains("http")
        (origin_list['protocol'].apply(lambda x: 'http' in x))
    ]
    # normal_df = origin_list
    if use_protocol:
        webshell_df = origin_list[
            # origin_list['protocol'].str.contains("http")
            (origin_list['protocol'].apply(lambda x: 'http' in x))
        ]

        dns_df = origin_list[
            (origin_list['protocol'].apply(lambda x: 'dns' in x))
        ]

        cs_df = origin_list[
            (origin_list['protocol'].apply(lambda x: 'dns' in x)) |
            (origin_list['protocol'].apply(lambda x: 'http' in x)) |
            (origin_list['protocol'].apply(lambda x: 'tls' in x)) |
            (origin_list['protocol'].apply(lambda x: 'icmp' in x))
            ]
    else:
        webshell_df = origin_list[
            (origin_list['source.bytes'] < 605245) &
            (origin_list['destination.bytes'] < 38889) &
            (origin_list['network.bytes'] < 644133) &
            (origin_list['protocol'].apply(lambda x: 'http' in x)) &
            (~origin_list['dns.host'].apply(lambda x: 'motor.groupama-sdig.com' in x)) &
            (~origin_list['dns.host'].apply(lambda x: 'image.groupama-sdig.com' in x))
            ]

        dns_df = origin_list[
            (origin_list['source.bytes'] < 240690) &
            (origin_list['destination.bytes'] < 45454) &
            (origin_list['network.bytes'] < 286144) &
            (origin_list['protocol'].apply(lambda x: 'dns' in x))
            ]

        cs_df = origin_list[
            (origin_list['source.bytes'] < 487) |
            (origin_list['destination.bytes'] < 427) |
            (origin_list['network.bytes'] < 914) |
            (origin_list['totDataBytes'] != 0)
            ]
    return normal_df.to_dict(orient='records'), webshell_df.to_dict(orient='records'), dns_df.to_dict(
        orient='records'), cs_df.to_dict(orient='records')


def clean_floats(data):
    if isinstance(data, dict):
        for key, value in data.items():
            data[key] = clean_floats(value)  # 递归处理字典中的数据
    elif isinstance(data, list):
        for i in range(len(data)):
            data[i] = clean_floats(data[i])  # 递归处理列表中的数据
    elif isinstance(data, decimal.Decimal):
        data = float(data)  # 转换 Decimal 为 float
    if isinstance(data, float):
        if math.isinf(data) or math.isnan(data):
            return -1
    return data


def list_of_dicts_to_csv(data):
    if not data:
        return ""
    headers = data[0].keys()
    csv_text = ",".join(headers) + "\n"
    for row in data:
        csv_text += ",".join(map(str, (row.get(header) for header in headers))) + "\n"
    return csv_text


def format_size(size):
    units = ["B", "KB", "MB", "GB", "TB"]
    unit_index = 0
    while size >= 1024 and unit_index < len(units) - 1:
        size /= 1024
        unit_index += 1
    return f"{size:.2f} {units[unit_index]}"


def get_total_size(obj, seen=None):
    """递归统计对象及其引用对象的总大小"""
    if seen is None:
        seen = set()
    obj_id = id(obj)
    if obj_id in seen:
        return 0  # 避免重复统计
    seen.add(obj_id)

    size = sys.getsizeof(obj)

    if isinstance(obj, dict):
        size += sum(get_total_size(k, seen) + get_total_size(v, seen) for k, v in obj.items())
    elif isinstance(obj, (list, tuple, set, frozenset)):
        size += sum(get_total_size(i, seen) for i in obj)

    return size


def clear_queue(save_queue):
    while True:
        try:
            save_queue.get_nowait()  # 尝试非阻塞地获取队列中的元素
        except queue.Empty:
            return
